package Share;

import java.io.*;
import java.math.BigInteger;
import java.net.*;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.SignatureException;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.X509EncodedKeySpec;

import javax.crypto.Cipher;
import javax.crypto.KeyAgreement;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.spec.DHParameterSpec;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

public class Session {
	private Socket socket; 
	private Certificate remoteCertificate;
	private Certificate localCertificate;
	private SecureDataOutputStream out;
	private SecureDataInputStream in;
	
	private static final BigInteger p = new BigInteger("fe76de0c179966141497196e0f5baee37550d3c42dd7b31159d78d1e966d54af5807fbcd094c8b7c643120fd35c2c4ef275ec8327fef8cc0babf83023c9a36f14ac9fdfb843df70c5a642cbe12e916f14db9c8d1c78195749a9845c49983efd528152ef020a5a6549c00c29fb181103f0b9a870c37a6a7efaff4e649affb7f43", 16);
	private static final BigInteger g = new BigInteger("868bb60bfa600cb0929105b78e5d30275ec45059aebaa9979c332f28268f691a0b74b4efa418d6df4866bc425ad6ff7f0870db0e6a4be1ad165fe009d11b6d4c28951f2d45f7b256a17e10100a9fa242623b454828de0facb2f4ca823f315c5ace46638d1cc754b4c4e24e698f49d8d473eb1d6d4b8182eaf2602756ada95e21", 16);
	private static final int l = 1023;
	
	public Certificate getRemoteCertificate()
	{
		return remoteCertificate;
	}
	
	public Certificate getLocalCertificate()
	{
		return localCertificate;
	}
	
	public Session(String server, int port, Certificate cert, PrivateKey key) throws UnknownHostException, IOException, SignatureException, InvalidKeyException, NoSuchAlgorithmException, InvalidAlgorithmParameterException, InvalidKeySpecException, NoSuchPaddingException
	{
		socket = new Socket(server, port);
		in = new SecureDataInputStream(new SecureInputStream(socket.getInputStream()));
		out = new SecureDataOutputStream(new SecureOutputStream(socket.getOutputStream()));
		socket.setSoTimeout(3600000);
		localCertificate = cert;
		setupCipher(cert, key);
	}
	
	public Session(Socket socket, Certificate cert, PrivateKey key) throws IOException, InvalidKeyException, SignatureException, InvalidKeySpecException, NoSuchAlgorithmException, InvalidAlgorithmParameterException, NoSuchPaddingException
	{
		this.socket = socket;
		in = new SecureDataInputStream(new SecureInputStream(socket.getInputStream()));
		out = new SecureDataOutputStream(new SecureOutputStream(socket.getOutputStream()));
		socket.setSoTimeout(3600000);
		localCertificate = cert;
		setupCipher(cert, key);
	}
	
	private void setupCipher(Certificate cert, PrivateKey key) throws IOException, InvalidKeyException, SignatureException, InvalidKeySpecException, NoSuchAlgorithmException, InvalidAlgorithmParameterException, NoSuchPaddingException
	{
		exchangeCertificate(cert);
		
		KeyPair dhKey = GenerateDHKey();
		
		byte[] dhPublicKey = dhKey.getPublic().getEncoded();
		Util.write(out, dhPublicKey);
		Util.write(out, Util.signData(dhPublicKey, key));
		
		byte[] remotePublicKey = Util.readByteArray(in);
		if(!Util.checkSig(remotePublicKey, Util.readByteArray(in), remoteCertificate))
		{
			throw new SignatureException();
		}
		SecretKeySpec spec = generateSecret(dhKey.getPrivate(), remotePublicKey);
		
		byte[] outIv = generateIV();
		Util.write(out, outIv);
		
		byte[] inIv = Util.readByteArray(in);
		
		Cipher cipherOut = generateCipher(Cipher.ENCRYPT_MODE, spec, outIv);
		Cipher cipherIn = generateCipher(Cipher.DECRYPT_MODE, spec, inIv);
		
		in.getStream().SetCipher(cipherIn);
		out.getStream().SetCipher(cipherOut);
	}
	
	private void exchangeCertificate(Certificate cert) throws InvalidKeyException, IOException
	{
		cert.serialize(out);
		remoteCertificate = new Certificate(in);
		if(!Util.verifyCertificate(remoteCertificate))
			throw new InvalidKeyException("The remote certificate was invalid");
	}
	
	private KeyPair GenerateDHKey() throws NoSuchAlgorithmException, InvalidAlgorithmParameterException
	{
		DHParameterSpec spec = new DHParameterSpec(p, g, l);
		KeyPairGenerator gen = KeyPairGenerator.getInstance("DH");
		gen.initialize(spec);
		return gen.genKeyPair();
	}
	
	private SecretKeySpec generateSecret(PrivateKey localPrivateKey, byte[] remoteKey) throws InvalidKeySpecException, NoSuchAlgorithmException, InvalidKeyException
	{
	    X509EncodedKeySpec x509KeySpec = new X509EncodedKeySpec(remoteKey);
	    KeyFactory keyFact = KeyFactory.getInstance("DH");
	    PublicKey remotePublicKey = keyFact.generatePublic(x509KeySpec);
		
	    KeyAgreement ka = KeyAgreement.getInstance("DH");
	    ka.init(localPrivateKey);
		ka.doPhase(remotePublicKey, true);
		
		
		byte[] secret = new byte[16];
		System.arraycopy(ka.generateSecret(), 0, secret, 0, secret.length);
		return new SecretKeySpec(secret,"AES");
	}
	
	private Cipher generateCipher(int opmode, SecretKeySpec spec, byte[] iv) throws InvalidKeyException, NoSuchAlgorithmException, InvalidKeySpecException, InvalidAlgorithmParameterException, NoSuchPaddingException
	{
		Cipher cipher = Cipher.getInstance("AES/CTR/NoPadding");
		cipher.init(opmode, spec, new IvParameterSpec(iv));
		return cipher;
	}
	
	public byte[] generateIV()
	{
		SecureRandom random = new SecureRandom();
		byte[] ivArray = new byte[16];
		random.nextBytes(ivArray);
		return ivArray;
	}
	
	public SecureDataInputStream getInputStream()
	{
		return in;
	}
	
	public boolean isClosed()
	{
		return socket.isClosed();
	}
	
	public SecureDataOutputStream getOutputStream()
	{
		return out;
	}
	
	public void close() throws IOException
	{
		socket.close();
	}
}
